﻿using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;

namespace Mars.GpuInstance
{
    [BurstCompile(FloatPrecision = FloatPrecision.Low, FloatMode = FloatMode.Fast)]
    public struct GPUSkinnedMeshAnimateJob : IJob
    {
        [ReadOnly]public NativeArray<float4> gpuiAnimationData;             // index: 0 x -> frameNo1, y -> frameNo2, z -> frameNo3, w -> frameNo4
                                                                            // index: 1 x -> weight1, y -> weight2, z -> weight3, w -> weight4

        [ReadOnly]public NativeArray<float4> gpuiAnimationTexture;         
        [ReadOnly]public int animationTextureSizeX;
        [ReadOnly]public int totalNumberOfBones;
        [ReadOnly]public int instanceCount;
        [WriteOnly]public NativeArray<float4x4> gpuiAnimationBuffer;

        public void Execute()
        {
            for (int i = 0; i < instanceCount; i++)
            {
                for (int j = 0; j < totalNumberOfBones; j++)
                {
                    var index = j + totalNumberOfBones * i;
                    gpuiAnimationBuffer[index] = CalculateBoneMatrix(i * 2, j);
                }
            }
        }
        
        float4x4 getBoneMatrixFromTexture(int frameIndex, int boneIndex)
        {
            // ordered by frames
            int textureIndex = frameIndex * totalNumberOfBones + boneIndex;

            int2 texIndex;
            texIndex.x = (textureIndex % animationTextureSizeX);
            texIndex.y = (int)(math.floor(textureIndex / (float) animationTextureSizeX) * 4);
            
            float4x4 boneMatrix = float4x4.identity;
            var pixelIndex = GetTextureIndex(texIndex);
            if (pixelIndex < 0 || pixelIndex >= gpuiAnimationTexture.Length)
            {
                return boneMatrix;
            }
            boneMatrix.c0 = gpuiAnimationTexture[pixelIndex];

            texIndex.y += 1;
            pixelIndex = GetTextureIndex(texIndex);
            if (pixelIndex < 0 || pixelIndex >= gpuiAnimationTexture.Length)
            {
                return boneMatrix;
            }
            boneMatrix.c1 = gpuiAnimationTexture[pixelIndex];
            
            
            texIndex.y += 1;
            pixelIndex = GetTextureIndex(texIndex);
            if (pixelIndex < 0 || pixelIndex >= gpuiAnimationTexture.Length)
            {
                return boneMatrix;
            }
            boneMatrix.c2 = gpuiAnimationTexture[pixelIndex];
            
            
            texIndex.y += 1;
            pixelIndex = GetTextureIndex(texIndex);
            if (pixelIndex < 0 || pixelIndex >= gpuiAnimationTexture.Length)
            {
                return boneMatrix;
            }
            boneMatrix.c3 = gpuiAnimationTexture[pixelIndex];
            return boneMatrix;
        }

        int GetTextureIndex(int2 texIndex)
        {
            int index = animationTextureSizeX * texIndex.y + texIndex.x;
            return index;
        }
        
        float4x4 CalculateBoneMatrix(int instanceIndex, int boneIndex)
        {
            float4 animationData = gpuiAnimationData[instanceIndex];
            float4 weightData = gpuiAnimationData[instanceIndex + 1];
    
            float frameNo = animationData[0];
            int frameStart = (int)math.floor(frameNo);
            frameStart = math.max(0, frameStart);

            float weight = weightData[0];
            float4x4 weightMatrix = new float4x4(
                weight, 0, 0, 0,
                0, weight, 0, 0,
                0, 0, weight, 0,
                0, 0, 0, weight);
            
            float4x4 bone = math.mul(getBoneMatrixFromTexture(frameStart, boneIndex), weightMatrix);
            // blending
            if (weight < 1)
            {
                for (int i = 1; i < 4; i++)
                {
                    weight = weightData[i];
                    if (weight > 0)
                    {
                        weightMatrix[0][0] = weight;
                        weightMatrix[1][1] = weight;
                        weightMatrix[2][2] = weight;
                        weightMatrix[3][3] = weight;

                        frameNo = animationData[i];
                        frameStart = (int)math.floor(frameNo);
                
                        bone += math.mul(getBoneMatrixFromTexture(frameStart, boneIndex), weightMatrix);
                    }
                }
            }

            return bone;
        }

        float4x4 lerp(float4x4 x, float4x4 y, float s)
        {
            float4x4 m = float4x4.identity;
            float4 l = new float4(s);
            m.c0 = math.lerp(x.c0, y.c0, l);
            m.c1 = math.lerp(x.c1, y.c1, l);
            m.c2 = math.lerp(x.c2, y.c2, l);
            m.c3 = math.lerp(x.c3, y.c3, l);
            return m;
        }
    }
}
